from enum import Enum
from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel
from app.schemas.metrics_schema import Chart, BarChart, ScatterPlot, LineChart
from app.schemas.assets_schema import Text, Ontology, KnowledgeGraph, GraphPattern,GraphRule ,MediaFile,TextRule
import numpy as np


# 1. 任务状态
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"


# 2. 产品类型
class ProductType(str, Enum):
    TEXT = "TEXT"
    FILE = "FILE"
    TEXT_RULE = "TEXT_RULE"
    GRAPH_RULE = "GRAPH_RULE"
    ONTOLOGY = "ONTOLOGY"
    KNOWLEDGE_GRAPH = "KNOWLEDGE_GRAPH"

# TSGE算法的输入参数
class TSGEInputParams(BaseModel):
    graph_sequence: List[Dict[str, Any]]  # 时序图序列
    embedding_dim: int = 128  # 嵌入维度
    learning_rate: float = 0.01  # 学习率
    epochs: int = 100  # 训练轮数
    node_types: Optional[Dict[str, str]] = None  # 节点类型信息

# TSGE算法的节点嵌入结果
class NodeEmbedding(BaseModel):
    node_id: str  # 节点ID
    embedding: List[float]  # 嵌入向量
    node_type: Optional[str] = None  # 节点类型

# TSGE算法的边预测结果
class LinkPrediction(BaseModel):
    source: str  # 源节点
    target: str  # 目标节点
    probability: float  # 连接概率
    time: Optional[float] = None  # 时间点

# TSGE算法的输出参数
class TSGEOutputParams(BaseModel):
    node_embeddings: List[NodeEmbedding]  # 节点嵌入结果
    link_predictions: Optional[List[LinkPrediction]] = None  # 边预测结果
    algorithm: str = "TSGE"  # 算法名称
    parameters: Dict[str, Any]  # 算法参数
    train_loss_history: Optional[List[float]] = None  # 训练损失历史

# 3. 输入参数
class InputParams(BaseModel):
    # 你的代码需要的输入参数
    tsge_params: Optional[TSGEInputParams] = None
    pass

# 4. 输出参数
class OutputParams(BaseModel):
    # 算法输出参数
    tsge_results: Optional[TSGEOutputParams] = None
    pass

# 5. 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams

# 6. 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: str
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    input_params: InputParams
    error_message: Optional[str] = None

    metrics: List[Chart] = []


# 7. 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None

    error_message: Optional[str] = None

    output_params: OutputParams

    metrics: List[Chart] = [] 